package com.dxysun.springboot.wrapper;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

public class ResponseWrapper extends HttpServletResponseWrapper
{

	private ByteArrayOutputStream buffer;
	private ServletOutputStream out;
	private PrintWriter writer;

	public ResponseWrapper(HttpServletResponse response) throws IOException
	{
		super(response);
		buffer = new ByteArrayOutputStream();
		out = new WrappedOutputStream(buffer);
		out.setWriteListener(new MyWriteListener());
		writer = new PrintWriter(new OutputStreamWriter(buffer, response.getCharacterEncoding()));
	}

	//重载父类获取OutputStream的方法
	@Override
	public ServletOutputStream getOutputStream() throws IOException
	{
		return out;
	}

	@Override
	public PrintWriter getWriter() throws IOException
	{
		return writer;
	}

	@Override
	public void flushBuffer() throws IOException
	{
		try
		{
			if (out != null)
			{
				out.flush();
				out.close();
			}
			if (writer != null)
			{
				writer.flush();
				writer.close();
			}
		}
		catch (IOException e)
		{
			e.printStackTrace();
		}
	}

	@Override
	public void reset()
	{
		buffer.reset();
	}

	public String getResponseData() throws IOException
	{
		flushBuffer();//将out、writer中的数据强制输出到WrappedResponse的buffer里面，否则取不到数据
		return buffer.toString(super.getCharacterEncoding());

	}

	public void writeResponseBody(String response) throws IOException
	{
//		super.getWriter().write(response);
//		super.getWriter().flush();
//		super.getWriter().close();

		super.getOutputStream().write(response.getBytes());
		super.getOutputStream().flush();
		super.getOutputStream().close();
	}

	//内部类，对ServletOutputStream进行包装，指定输出流的输出端
	private class WrappedOutputStream extends ServletOutputStream
	{

		private ByteArrayOutputStream bos;

		public WrappedOutputStream(ByteArrayOutputStream stream) throws IOException
		{
			bos = stream;
		}

		//将指定字节写入输出流bos
		@Override
		public void write(int b) throws IOException
		{
			bos.write(b);
		}

		@Override
		public boolean isReady()
		{
			return false;
		}

		@Override
		public void setWriteListener(WriteListener listener)
		{

		}
	}

	private class MyWriteListener implements WriteListener
	{

		@Override
		public void onWritePossible() throws IOException {
			System.out.println("HttpServletResponseWrapper onWritePossible");
		}

		@Override
		public void onError(Throwable throwable) {
			System.out.println("HttpServletResponseWrapper onError");

		}
	}

}
